#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Jun  8 14:48:38 2025

@author: fenghongli
"""

import tushare as ts
import pandas as pd
import os

ts.set_token('2b46ae1723755362c4a7b5bf50d174b4469669fea505758c9b1853cf')
pro = ts.pro_api()

# 创建数据保存目录
output_dir = 'stock_returns'
os.makedirs(output_dir, exist_ok=True)

# 设置时间区间
start_date = '20110101'
end_date = '20241231'

# 股票代码列表
# 银行、保险、证券（金融类20家）
financial_stocks = [
    '600036.SH',  # 招商银行
    '601398.SH',  # 工商银行
    '601939.SH',  # 建设银行
    '601288.SH',  # 农业银行
    '600000.SH',  # 浦发银行
    '601988.SH',  # 中国银行
    '601328.SH',  # 交通银行
    '600016.SH',  # 民生银行
    '601336.SH',  # 新华保险
    '601628.SH',  # 中国人寿
    '601601.SH',  # 中国太保
    '601318.SH',  # 中国平安
    '600030.SH',  # 中信证券
    '601688.SH',  # 华泰证券
    '600837.SH',  # 海通证券
    '000001.SZ',  # 平安银行
    '000776.SZ',  # 广发证券
    '002142.SZ',  # 宁波银行
    '601577.SH',  # 长沙银行
    '600999.SH',  # 招商证券
]

# 房地产类10家（头部上市房企）
real_estate_stocks = [
    '000002.SZ',  # 万科A
    '600048.SH',  # 保利发展
    '600340.SH',  # 华夏幸福
    '600383.SH',  # 金地集团
    '000656.SZ',  # 金科股份
    '000671.SZ',  # 阳光城
    '000069.SZ',  # 华侨城
    '600606.SH',  # 绿地控股
    '600325.SH',  # 华发股份
    '600223.SH',  # 鲁商发展
]

all_stocks = financial_stocks + real_estate_stocks

# 循环获取数据并计算收益率
for ts_code in all_stocks:
    try:
        print(f"正在获取 {ts_code} 的数据...")
        df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
        df.sort_values('trade_date', inplace=True)
        df['ret'] = df['close'].pct_change()  # 计算日收益率

        # 保存数据
        save_name = f"{ts_code.replace('.', '_')}_returns.csv"
        df[['trade_date', 'ts_code', 'close', 'ret']].to_csv(os.path.join(output_dir, save_name), index=False)
    except Exception as e:
        print(f"获取 {ts_code} 时出错：{e}")

print("所有数据已完成提取并保存到 'stock_returns/' 文件夹。")
